K-th largest element in an array

Time: O(N) ~ O(N^2); Space: O(1); medium

Find the k-th largest element in an unsorted array.

Note:

  • It is the k-th largest element in the sorted order, not the k-th distinct element.

Example 1:

Input: nums = [3,2,1,5,6,4], k = 2

Output: 5

Example 2:

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4

Output: 4

[1]:

from random import randint

class Solution1(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: {int[]}
        :type k: int
        :rtype: int
        """
        left, right = 0, len(nums) - 1

        while left <= right:
            pivot_idx = randint(left, right)
            new_pivot_idx = self.PartitionAroundPivot(left, right, pivot_idx, nums)

            if new_pivot_idx == k - 1:
                return nums[new_pivot_idx]
            elif new_pivot_idx > k - 1:
                right = new_pivot_idx - 1
            else:                    # new_pivot_idx < k - 1
                left = new_pivot_idx + 1

    def PartitionAroundPivot(self, left, right, pivot_idx, nums):
        pivot_value = nums[pivot_idx]
        new_pivot_idx = left
        nums[pivot_idx], nums[right] = nums[right], nums[pivot_idx]

        for i in range(left, right):
            if nums[i] > pivot_value:
                nums[i], nums[new_pivot_idx] = nums[new_pivot_idx], nums[i]
                new_pivot_idx += 1

        nums[right], nums[new_pivot_idx] = nums[new_pivot_idx], nums[right]

        return new_pivot_idx

[2]:

s = Solution1()
nums = [3,2,1,5,6,4]
k = 2
assert s.findKthLargest(nums, k) == 5
nums = [3,2,3,1,2,4,5,5,6]
k = 4
assert s.findKthLargest(nums, k) == 4